package com.shujia;

import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.StructField;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.io.Text;
import org.apache.tools.ant.taskdefs.Length;

import java.util.ArrayList;
import java.util.List;

/**
 * hello world hello
 * 返回三行：
 * hello
 * world
 * hello
 * <p>
 * 样例：
 * select mysplit("hello world hello"," ")
 * <p>
 * add jar /usr/local/soft/jars/hiveCode-1.0.jar
 * create temporary function mysplit as 'com.shujia.MyUDTF';
 */
public class MyUDTF extends GenericUDTF {

    @Override
    public StructObjectInspector initialize(StructObjectInspector argOIs) throws UDFArgumentLengthException, UDFArgumentTypeException {
        // 通过StructObjectInspector获取所有参数，其中StructField表示对应一个参数
        List<? extends StructField> allStructFieldRefs = argOIs.getAllStructFieldRefs();

        // 判断参数个数
        if (allStructFieldRefs.size() != 2) {
            throw new UDFArgumentLengthException("该函数参数为两个，其中一个为待分割的字符串，另一个为分隔符");
        }

        // 判断参数类型
        for (int i = 0; i < allStructFieldRefs.size(); i++) {
            StructField structField = allStructFieldRefs.get(i);
            ObjectInspector objectInspector = structField.getFieldObjectInspector();
            if (!objectInspector.getCategory().equals(ObjectInspector.Category.PRIMITIVE)) {
                throw new UDFArgumentTypeException(i + 1, "参数的类型不正确，请使用HIVE中的普通数据类型String");
            }
        }

        // 返回字段名称
        ArrayList<String> fieldNames = new ArrayList<String>();
        fieldNames.add("word");


        // 返回字段的类型
        ArrayList<ObjectInspector> fieldOIs = new ArrayList<ObjectInspector>();
//        fieldOIs.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector);

        // Hadoop中的序列化数据类型writableString-> Hadoop中的Text
        fieldOIs.add(PrimitiveObjectInspectorFactory.writableStringObjectInspector);


        return ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames,
                fieldOIs);

    }


    @Override
    public void process(Object[] args) throws HiveException {
        String line_words = args[0].toString();
        String split = args[1].toString();
        String[] words = line_words.split(split);

        // String[] 表示为一行数据，由于最终返回列只有一个并且类型为JavaString，所以可以使用String[]
//        String[] returnRes = new String[1];

        //
        ArrayList<Text> returnRes = new ArrayList<>();

        for (String word : words) {
            // 将切分好的单词，通过String[]进行保存，然后通过forward函数返回出去
//            returnRes[0] = word;

            returnRes.add(new Text(word));
            forward(returnRes);
        }
    }

    @Override
    public void close() throws HiveException {
        // TODO: 2022/9/15  nothing to do
    }
}
